#include <quadrotor_msgs/PositionCommand.h>
// #include <kr_mav_msgs/TRPYCommand.h>
#include <std_msgs/Float32MultiArray.h>
#include <nav_msgs/Odometry.h>
#include <nodelet/nodelet.h>
#include <ros/ros.h>
#include <std_msgs/Bool.h>
#include <tf/transform_datatypes.h>
#include <uav_utils/geometry_utils.h>

#include <Eigen/Geometry>
#include <pid_control/PIDControl.hpp>

#define CLAMP(x, min, max) ((x) < (min)) ? (min) : ((x) > (max)) ? (max) : (x)

class PIDControlNodelet : public nodelet::Nodelet
{
 public:
  PIDControlNodelet()
      : position_cmd_updated_(false),
        position_cmd_init_(false),
        des_yaw_(0),
        current_yaw_(0),
        enable_motors_(false),
        use_external_yaw_(false)
  {
  }

  void onInit(void);

 private:
  void publishTRPYCommand(void);
  void publishRPMCommand(void);
  void position_cmd_callback(const quadrotor_msgs::PositionCommand::ConstPtr &cmd);
  void odom_callback(const nav_msgs::Odometry::ConstPtr &odom);
  void enable_motors_callback(const std_msgs::Bool::ConstPtr &msg);

  PIDControl controller_;
  ros::Publisher trpy_command_pub_;
  ros::Publisher  rpm_command_pub_;

  ros::Subscriber odom_sub_;
  ros::Subscriber position_cmd_sub_;
  ros::Subscriber enable_motors_sub_;

  bool position_cmd_updated_, position_cmd_init_;
  std::string frame_id_;

  Eigen::Vector3f des_pos_, des_vel_, des_acc_, kx_, kv_, ki_, kR_, kOm_;
  float          init_x_, init_y_, init_z_;
  float des_yaw_;
  float current_yaw_;
  float ki_yaw_;
  float max_roll_pitch_;
  bool enable_motors_;
  bool use_external_yaw_;
  Eigen::Matrix3d rotation_matrix;
  Eigen::Vector3d omega;
};

void PIDControlNodelet::publishRPMCommand(void)
{
  // cal trpy cmd first
  float ki_yaw = 0;
  Eigen::Vector3f ki = Eigen::Vector3f::Zero();
  // Only enable integral terms when motors are on
  if(1)
  {
    ki_yaw = ki_yaw_;
    ki = ki_;
  }
  controller_.calculateControl(des_pos_, des_vel_, des_acc_, des_yaw_, kx_, kv_, ki, ki_yaw);

  const Eigen::Vector4f &trpy = controller_.getControls();

  double thrust{0.0}, roll{0.0}, pitch{0.0}, yaw{0.0};

  if(1)
  {
    thrust = trpy(0);
    roll = CLAMP(trpy(1), -max_roll_pitch_, max_roll_pitch_);
    pitch = CLAMP(trpy(2), -max_roll_pitch_, max_roll_pitch_);
    yaw = trpy(3);
  }
  
  // convert trpy cmd to actuator controls
  // TODO: add model paramers
  const double _kf = 8.98132e-9;
  const double _km = 0.07 * (3 * 0.062) * _kf;
  const double kf = _kf;
  const double km = _km / _kf * kf;

  const double d = 0.26;
  float Ixx = 2.64e-3, Iyy = 2.64e-3, Izz = 4.96e-3;
  const Eigen::Matrix3f J = Eigen::Vector3f(Ixx, Iyy, Izz).asDiagonal();

  const float           I[3][3] = { { J(0, 0), J(0, 1), J(0, 2) },
                          { J(1, 0), J(1, 1), J(1, 2) },
                          { J(2, 0), J(2, 1), J(2, 2) } };
  
  // Rotation, may use external yaw
  Eigen::Vector3d _ypr = uav_utils::R_to_ypr(rotation_matrix);
  Eigen::Vector3d ypr  = _ypr;

  if (use_external_yaw_)
    ypr[0] = current_yaw_;

  Eigen::Matrix3d R;
  R = Eigen::AngleAxisd(ypr[0], Eigen::Vector3d::UnitZ()) *
      Eigen::AngleAxisd(ypr[1], Eigen::Vector3d::UnitY()) *
      Eigen::AngleAxisd(ypr[2], Eigen::Vector3d::UnitX());

  float R11 = R(0, 0);
  float R12 = R(0, 1);
  float R13 = R(0, 2);
  float R21 = R(1, 0);
  float R22 = R(1, 1);
  float R23 = R(1, 2);
  float R31 = R(2, 0);
  float R32 = R(2, 1);
  float R33 = R(2, 2);

  float Om1 = omega(0);
  float Om2 = omega(1);
  float Om3 = omega(2);

  float Rd11 = cos(yaw) * cos(pitch);
  float Rd12 = cos(yaw) * sin(pitch) * sin(roll) - cos(roll) * sin(yaw);
  float Rd13 = sin(yaw) * sin(roll) + cos(yaw) * cos(roll) * sin(pitch);
  float Rd21 = cos(pitch) * sin(yaw);
  float Rd22 = cos(yaw) * cos(roll) + sin(yaw) * sin(pitch) * sin(roll);
  float Rd23 = cos(roll) * sin(yaw) * sin(pitch) - cos(yaw) * sin(roll);
  float Rd31 = -sin(pitch);
  float Rd32 = cos(pitch) * sin(roll);
  float Rd33 = cos(pitch) * cos(roll);

  float Psi = 0.5f * (3.0f - (Rd11 * R11 + Rd21 * R21 + Rd31 * R31 + Rd12 * R12 + Rd22 * R22 + Rd32 * R32 + Rd13 * R13 +
                              Rd23 * R23 + Rd33 * R33));

  float force = 0;
  if(Psi < 1.0f)  // Position control stability guaranteed only when Psi < 1
    force = thrust;

  float eR1 = 0.5f * (R12 * Rd13 - R13 * Rd12 + R22 * Rd23 - R23 * Rd22 + R32 * Rd33 - R33 * Rd32);
  float eR2 = 0.5f * (R13 * Rd11 - R11 * Rd13 - R21 * Rd23 + R23 * Rd21 - R31 * Rd33 + R33 * Rd31);
  float eR3 = 0.5f * (R11 * Rd12 - R12 * Rd11 + R21 * Rd22 - R22 * Rd21 + R31 * Rd32 - R32 * Rd31);
  
  // TODO: angular_velocity feedforward
  float angular_velocity[3] = {0.0f, 0.0f, 0.0f};

  float Omd1 = angular_velocity[0] * (R11 * Rd11 + R21 * Rd21 + R31 * Rd31) +
               angular_velocity[1] * (R11 * Rd12 + R21 * Rd22 + R31 * Rd32) +
               angular_velocity[2] * (R11 * Rd13 + R21 * Rd23 + R31 * Rd33);
  float Omd2 = angular_velocity[0] * (R12 * Rd11 + R22 * Rd21 + R32 * Rd31) +
               angular_velocity[1] * (R12 * Rd12 + R22 * Rd22 + R32 * Rd32) +
               angular_velocity[2] * (R12 * Rd13 + R22 * Rd23 + R32 * Rd33);
  float Omd3 = angular_velocity[0] * (R13 * Rd11 + R23 * Rd21 + R33 * Rd31) +
               angular_velocity[1] * (R13 * Rd12 + R23 * Rd22 + R33 * Rd32) +
               angular_velocity[2] * (R13 * Rd13 + R23 * Rd23 + R33 * Rd33);

  float eOm1 = Om1 - Omd1;
  float eOm2 = Om2 - Omd2;
  float eOm3 = Om3 - Omd3;

#if 0
  float in1 = Om2 * (I[2][0] * Om1 + I[2][1] * Om2 + I[2][2] * Om3) -
              Om3 * (I[1][0] * Om1 + I[1][1] * Om2 + I[1][2] * Om3);
  float in2 = Om3 * (I[0][0] * Om1 + I[0][1] * Om2 + I[0][2] * Om3) -
              Om1 * (I[2][0] * Om1 + I[2][1] * Om2 + I[2][2] * Om3);
  float in3 = Om1 * (I[1][0] * Om1 + I[1][1] * Om2 + I[1][2] * Om3) -
              Om2 * (I[0][0] * Om1 + I[0][1] * Om2 + I[0][2] * Om3);
#else
  float in1 = Omd2 * (I[2][0] * Omd1 + I[2][1] * Omd2 + I[2][2] * Omd3) -
              Omd3 * (I[1][0] * Omd1 + I[1][1] * Omd2 + I[1][2] * Omd3);
  float in2 = Omd3 * (I[0][0] * Omd1 + I[0][1] * Omd2 + I[0][2] * Omd3) -
              Omd1 * (I[2][0] * Omd1 + I[2][1] * Omd2 + I[2][2] * Omd3);
  float in3 = Omd1 * (I[1][0] * Omd1 + I[1][1] * Omd2 + I[1][2] * Omd3) -
              Omd2 * (I[0][0] * Omd1 + I[0][1] * Omd2 + I[0][2] * Omd3);
#endif

  float M1 = -kR_[0] * eR1 - kOm_[0] * eOm1 + in1;
  float M2 = -kR_[1] * eR2 - kOm_[1] * eOm2 + in2;
  float M3 = -kR_[2] * eR3 - kOm_[2] * eOm3 + in3;
  
  // mixers
  float w_sq[4];
  w_sq[0] = force / (4 * kf) - M1 / (4 * d * kf) - M2 / (4 * d * kf) - M3 / (4 * km);
  w_sq[1] = force / (4 * kf) + M1 / (4 * d * kf) + M2 / (4 * d * kf) - M3 / (4 * km);
  w_sq[2] = force / (4 * kf) + M1 / (4 * d * kf) - M2 / (4 * d * kf) + M3 / (4 * km);
  w_sq[3] = force / (4 * kf) - M1 / (4 * d * kf) + M2 / (4 * d * kf) + M3 / (4 * km);


  std_msgs::Float32MultiArray rpm_array;
  float rpm_data[4];
  rpm_array.data.clear();

  for (int i = 0; i < 4; i++)
  {
    if (w_sq[i] < 0)
      w_sq[i] = 0;

    rpm_data[i] = sqrtf(w_sq[i]);
    rpm_array.data.push_back(rpm_data[i]);
  }

  rpm_command_pub_.publish(rpm_array);
}

// void PIDControlNodelet::publishTRPYCommand(void)
// {
//   float ki_yaw = 0;
//   Eigen::Vector3f ki = Eigen::Vector3f::Zero();
//   // Only enable integral terms when motors are on
//   if(enable_motors_)
//   {
//     ki_yaw = ki_yaw_;
//     ki = ki_;
//   }
//   controller_.calculateControl(des_pos_, des_vel_, des_acc_, des_yaw_, kx_, kv_, ki, ki_yaw);

//   const Eigen::Vector4f &trpy = controller_.getControls();

//   kr_mav_msgs::TRPYCommand::Ptr trpy_command(new kr_mav_msgs::TRPYCommand);
//   trpy_command->header.stamp = ros::Time::now();
//   trpy_command->header.frame_id = frame_id_;
//   if(enable_motors_)
//   {
//     trpy_command->thrust = trpy(0);
//     trpy_command->roll = CLAMP(trpy(1), -max_roll_pitch_, max_roll_pitch_);
//     trpy_command->pitch = CLAMP(trpy(2), -max_roll_pitch_, max_roll_pitch_);
//     trpy_command->yaw = trpy(3);
//   }
//   trpy_command->aux.current_yaw = current_yaw_;
//   trpy_command->aux.enable_motors = enable_motors_;
//   trpy_command->aux.use_external_yaw = use_external_yaw_;
//   trpy_command_pub_.publish(trpy_command);
// }

void PIDControlNodelet::position_cmd_callback(const quadrotor_msgs::PositionCommand::ConstPtr &cmd)
{
  des_pos_ = Eigen::Vector3f(cmd->position.x, cmd->position.y, cmd->position.z);
  des_vel_ = Eigen::Vector3f(cmd->velocity.x, cmd->velocity.y, cmd->velocity.z);
  des_acc_ = Eigen::Vector3f(cmd->acceleration.x, cmd->acceleration.y, cmd->acceleration.z);
  // kx_ = Eigen::Vector3f(cmd->kx[0], cmd->kx[1], cmd->kx[2]);
  // kv_ = Eigen::Vector3f(cmd->kv[0], cmd->kv[1], cmd->kv[2]);

  des_yaw_ = cmd->yaw;
  position_cmd_updated_ = true;
  position_cmd_init_ = true;

  // publishTRPYCommand();
  publishRPMCommand();
}

void PIDControlNodelet::odom_callback(const nav_msgs::Odometry::ConstPtr &odom)
{
  const Eigen::Vector3f position(odom->pose.pose.position.x, odom->pose.pose.position.y, odom->pose.pose.position.z);
  const Eigen::Vector3f velocity(odom->twist.twist.linear.x, odom->twist.twist.linear.y, odom->twist.twist.linear.z);

  const Eigen::Quaterniond quat(odom->pose.pose.orientation.w, odom->pose.pose.orientation.x, 
                                odom->pose.pose.orientation.y, odom->pose.pose.orientation.z);
  
  const Eigen::Vector3d _omega(odom->twist.twist.angular.x,
                              odom->twist.twist.angular.y,
                              odom->twist.twist.angular.z);
  omega = _omega; 

  rotation_matrix = quat.toRotationMatrix();

  current_yaw_ = tf::getYaw(odom->pose.pose.orientation);

  controller_.setPosition(position);
  controller_.setVelocity(velocity);
  controller_.setYaw(current_yaw_);

  if(position_cmd_init_)
  {
    // We set position_cmd_updated_ = false and expect that the
    // position_cmd_callback would set it to true since typically a position_cmd
    // message would follow an odom message. If not, the position_cmd_callback
    // hasn't been called and we publish the so3 command ourselves
    // TODO: Fallback to hover if position_cmd hasn't been received for some time
    if(!position_cmd_updated_)
      // publishTRPYCommand();
      publishRPMCommand();
    position_cmd_updated_ = false;
  }
  else if (init_z_ > -9999.0)
  {
    des_pos_ = Eigen::Vector3f(init_x_, init_y_, init_z_);
    des_vel_ = Eigen::Vector3f(0,0,0);
    des_acc_ = Eigen::Vector3f(0,0,0);

    publishRPMCommand();
  }
  
}

void PIDControlNodelet::enable_motors_callback(const std_msgs::Bool::ConstPtr &msg)
{
  if(msg->data)
    ROS_INFO("Enabling motors");
  else
    ROS_INFO("Disabling motors");

  enable_motors_ = msg->data;

  // Reset integrals when toggling motor state
  controller_.resetIntegrals();
}

void PIDControlNodelet::onInit(void)
{
  ros::NodeHandle n(getPrivateNodeHandle());

  std::string quadrotor_name;
  n.param("quadrotor_name", quadrotor_name, std::string("quadrotor"));
  frame_id_ = "/" + quadrotor_name;

  double mass;
  n.param("mass", mass, 0.235);
  controller_.setMass(mass);
  controller_.setMaxIntegral(mass * 3);

  n.param("use_external_yaw", use_external_yaw_, true);

  double max_roll_pitch;
  n.param("max_roll_pitch", max_roll_pitch, 30.0);
  max_roll_pitch_ = max_roll_pitch;

  double ki_x, ki_y, ki_z, ki_yaw;
  n.param("gains/rot/x", kR_[0], 1.5f);
  n.param("gains/rot/y", kR_[1], 1.5f);
  n.param("gains/rot/z", kR_[2], 1.0f);
  n.param("gains/ang/x", kOm_[0], 0.13f);
  n.param("gains/ang/y", kOm_[1], 0.13f);
  n.param("gains/ang/z", kOm_[2], 0.1f);
  n.param("gains/kx/x", kx_[0], 5.7f);
  n.param("gains/kx/y", kx_[1], 5.7f);
  n.param("gains/kx/z", kx_[2], 6.2f);
  n.param("gains/kv/x", kv_[0], 3.4f);
  n.param("gains/kv/y", kv_[1], 3.4f);
  n.param("gains/kv/z", kv_[2], 4.0f);
  n.param("gains/ki/x", ki_x, 0.0);
  n.param("gains/ki/y", ki_y, 0.0);
  n.param("gains/ki/z", ki_z, 0.0);
  n.param("gains/ki/yaw", ki_yaw, 0.01);
  ki_[0] = ki_x, ki_[1] = ki_y, ki_[2] = ki_z;
  ki_yaw_ = ki_yaw;

  controller_.resetIntegrals();

  n.param("pid_control/init_state_x", init_x_, 0.0f);
  n.param("pid_control/init_state_y", init_y_, 0.0f);
  n.param("pid_control/init_state_z", init_z_, -10000.0f);

  // trpy_command_pub_ = n.advertise<kr_mav_msgs::TRPYCommand>("trpy_cmd", 10);
  rpm_command_pub_ = n.advertise<std_msgs::Float32MultiArray>("rpm_cmd", 10);


  odom_sub_ = n.subscribe("odom", 10, &PIDControlNodelet::odom_callback, this, ros::TransportHints().tcpNoDelay());
  position_cmd_sub_ = n.subscribe("position_cmd", 10, &PIDControlNodelet::position_cmd_callback, this,
                                  ros::TransportHints().tcpNoDelay());

  enable_motors_sub_ =
      n.subscribe("motors", 2, &PIDControlNodelet::enable_motors_callback, this, ros::TransportHints().tcpNoDelay());
}

#include <pluginlib/class_list_macros.h>
PLUGINLIB_EXPORT_CLASS(PIDControlNodelet, nodelet::Nodelet);
